#!/usr/bin/env python

import gc
import socket
import time
import itertools
import sys
import getopt
import signal
import os
import errno
from math import floor

if sys.platform == "win32":
    timer = time.clock
else:
    timer = time.time

def debug( level, * msgs ):
    if debug.level >= level:
        for m in msgs:
            print >> sys.stderr, m,
        print >> sys.stderr
# debug()
debug.level = 0


class StressTest( object ):
    """Plataform to stress test servers.

    This plataform is designed to help stress test bank servers, it has the
    facility to provide a new request id and the account to operate on.

    Recommended usage is to inherit from this class and override the
    C{action()} method. If you want some setup, you can also override
    C{setup()}.

    To send commands please use L{send_command()<StressTest.send_command>},
    to receive the acknowledge reply, use L{recv_ok()<StressTest.recv_ok>}.

    You may use C{self.socket} to operate directly over the socket or the
    easier to use C{self.rfile} (input) and C{self.wfile} (output), so you
    get operations like readline.

    """
    request_id = 0
    account    = 0

    def __init__( self, server_address ):
        """Creates a new StressTest that will be connected to server address.

        @param server_address: a tuple C{( address, port )} to connect to:
                               C{( 'localhost', 8000 )}
        @type  server_address: L{tuple}
        """
        self.socket = socket.socket( socket.AF_INET, socket.SOCK_STREAM )
        self.socket.connect( server_address )
        self.rfile = self.socket.makefile( 'r' )
        self.wfile = self.socket.makefile( 'w' )
    # __init__()


    def next_request_id( classobj ):
        """Increment internal request ID.

        @note: the request_id is shared among all StressTest descendant, so
        it's safe to call multiple tests, the request_id will not be an already
        used one.
        """
        classobj.request_id += 1
    # inc_counter()
    next_request_id = classmethod( next_request_id )


    def set_request_id( classobj, request_id ):
        """Set request_id to new value.

        @param request_id: new value
        @type  request_id: L{int}
        """
        if not isinstance( request_id, int ):
            raise TypeError( "request_id must be an integer!" )
        classobj.request_id = request_id
    # set_request_id()
    set_request_id = classmethod( set_request_id )


    def next_account( classobj ):
        """Increment internal account.

        @note: the account is shared among all StressTest descendant, so it's
        safe to call multiple tests, the account will not be an already used
        one.
        """
        classobj.account += 1
    # inc_counter()
    next_account = classmethod( next_account )


    def send( self, value ):
        """Send value over the socket.

        @param value: what to send over the socket.
        """
        count = 0
        maxcount = 3
        while count < maxcount:
            try:
                self.wfile.write( value )
                self.wfile.flush()
                return
            except socket.error, e:
                if e[ 0 ] != errno.EINTR:
                    debug( 1, "socket error: %s" % e )
                    debug( 1, "Trying again... %s/%s" % ( count, maxcount ) )
                count += 1
    # send()
    write = send


    def send_command( self, command, *arguments ):
        """Send a command over the socket.

        This will use the request_id and call
        L{next_request_id()<StressTest.next_request_id>} after it.

        @param command: command name
        @param arguments: arguments to command.
        """
        cmd = "%s %d %s\n" % ( command,
                               self.request_id,
                               " ".join( [ str( a ) for a in arguments ] ) )
        debug( 3, "send:", cmd )
        self.send( cmd )
        self.next_request_id()
    # send_command()


    def recv( self, bytes ):
        """Receive C{bytes} from socket.

        @return: read bytes from socket.
        @rtype: L{str}
        """
        count = 0
        maxcount = 3
        while count < maxcount:
            try:
                v = self.rfile.read( bytes )
                debug( 3, "recv:", repr( v ) )
                return v
            except socket.error, e:
                if e[ 0 ] != errno.EINTR:
                    debug( 1, "socket error: %s" % e )
                    debug( 1, "Trying again... %s/%s" % ( count, maxcount ) )
                count += 1
    # recv()
    read = recv


    def recv_line( self ):
        """Receive a line (with trailing \n).

        @return: the read line.
        @rtype: L{str}
        """
        v = self.rfile.readline()
        debug( 3, "recv:", repr( v ) )
        return v
    # recv_line()
    readline = recv_line


    def recv_reply( self ):
        """Receive reply to command.

        The protocol states that a reply is a bunch of lines and ends at '\n\n'
        so we return everything but these two new line.

        @return: reply to command
        @rtype: L{str}
        """
        lines = []
        while True:
            l = self.readline()
            lines.append( l[ : -1 ] ) # chomp \n
            if len( lines ) > 1 and l == '\n':
                break

        del lines[ -1 ] # remove the last empty line
        return "\n".join( lines )
    # recv_reply()


    def recv_ok( self ):
        """Receive the acknowledgment from server.

        @return: True if received the ack, False otherwise.
        @rtype: L{bool}
        """
        ok_str = "+OK\n\n"
        v =  self.recv( len( ok_str ) )
        if v != ok_str:
            debug( 1, "Could not receive %r, got %r instead." % ( ok_str, v ) )
            return False
        else:
            return True
    # recv_ok()


    def setup( self ):
        """Setup before running tests.

        @note: this may be overridden, but remember to call the parent's
               method.
        """
        self.next_account()
    # setup()


    def action( self ):
        """Perform the test action.

        This must be overridden to provide the test functionality.
        """
        pass
    # action()


    def check( self, times ):
        """Check if operation was successfully performed C{times}.

        This must be overridden to provide the check if everything worked all
        right.

        It should return True if everythin went all right or raise a
        ValueError with a meaningful message.
        """
        raise Warning( "Tests must override check()!" )
        return True
    # check()


    def timeit( self, func, times ):
        """Times how long it takes to execute C{times} C{func()}.

        @return: time in seconds.
        @rtype: L{float}
        """
        before = timer()
        last   = 0.0
        last_t = before
        for i in xrange( times ):
            v = floor( i * 10.0 / times )
            if v > last:
                now = timer()
                debug( 1, "done: %s%%, last 10%% in %s seconds." % \
                       ( int( v * 10 ), now - last_t ) )
                last = v
                last_t = now
            func()

        if not self.check( times ):
            debug( 2, "Checking for operations returned False but no "
                    "exception was raised, using default ValueError." )
            raise ValueError( ( "executing %s this operation was not " \
                                "successful" ) % times )
        else:
            debug( 2, "%s operations executed as expected." % times )

        after = timer()
        return after - before
    # timeit()


    def run( self, times=1000 ):
        """Runs C{self.action()} C{times} and returns statistics about it.

        @return: a tuple in format:
                 C{( number of seconds, operations per second)}
        @rtype: L{tuple} of L{float}s
        """
        self.setup()

        try:
            v = self.timeit( self.action, times )
        except KeyboardInterrupt:
            debug( 0, "User required program to stop using keyboard" )

        ret = ( v, float( times )/ v )
        return ret
    # run()


    def howmany( self, maxtime=60 ):
        """Gather how many C{self.action()} calls can be made in C{maxtime}
        seconds.

        @return: a tuple in format:
                 C{( number of calls, total time, seconds per operation )}
        @rtype: L{tuple} of L{float}s
        """
        self.setup()

        try:
            _t = timer # avoid many name lookups
            start   = _t()
            now     = _t()
            action  = self.action
            counter = 0 # how many loops counter
            while now - start < maxtime:
                action()
                counter += 1
                now = _t()

            if not self.check( counter ):
                debug( 2, "Checking for operations returned False but no "
                       "exception was raised, using default ValueError." )
                raise ValueError( ( "executing %s this operation was not " \
                                    "successful" ) % counter )
            else:
                debug( 2, "%s operations executed as expected." % counter )

        except KeyboardInterrupt:
            debug( 0, "User required program to stop using keyboard" )

        realtime = now - start
        return ( counter, realtime, realtime / counter )
    # howmany()
# StressTest


class StressTestDeposit( StressTest ):
    """Stress Test 'deposit' operation."""
    def setup( self ):
        StressTest.setup( self )
        self.send_command( "consult", self.account )
        assert( self.recv_ok() )
        self.base_val = int( self.recv_reply() )
    # setup()


    def action( self ):
        self.send_command( "deposit", self.account, 10 )
        assert( self.recv_ok() )
    # action()


    def check( self, times ):
        self.send_command( "consult", self.account )
        assert( self.recv_ok() )
        v = int( self.recv_reply() )
        if v != self.base_val + times * 10:
            raise ValueError( ( "After %s deposits of %s with initial "
                                "value %s, got %s instead of the correct "
                                "value %s." ) % \
                              ( times, 10, self.base_val, v,
                                self.base_val + times * 10 ) )
        else:
            return True
    # check()
# StressTestDeposit


class StressTestWithdraw( StressTest ):
    """Stress Test 'withdraw' operation."""
    def setup( self ):
        StressTest.setup( self )
        self.send_command( "consult", self.account )
        assert( self.recv_ok() )
        self.base_val = int( self.recv_reply() )
    # setup()


    def action( self ):
        self.send_command( "withdraw", self.account, 10 )
        assert( self.recv_ok() )
    # action()


    def check( self, times ):
        self.send_command( "consult", self.account )
        assert( self.recv_ok() )
        v = int( self.recv_reply() )
        if v != self.base_val - times * 10:
            raise ValueError( ( "After %s withdraws of %s with initial "
                                "value %s, got %s instead of the correct "
                                "value %s." ) % \
                              ( times, 10, self.base_val, v,
                                self.base_val - times * 10 ) )
        else:
            return True
    # check()
# StressTestWithdraw


class StressTestTransfer( StressTest ):
    """Stress Test 'transfer' operation."""
    def setup( self ):
        StressTest.setup( self )
        self.src = self.account
        self.next_account()
        self.dst = self.account

        self.send_command( "consult", self.src )
        assert( self.recv_ok() )
        self.src_base_val = int( self.recv_reply() )
        self.send_command( "consult", self.dst )
        assert( self.recv_ok() )
        self.dst_base_val = int( self.recv_reply() )
    # setup()

    def action( self ):
        self.send_command( "transfer", self.src, self.dst, 10 )
        assert( self.recv_ok() == True )
    # action()


    def check( self, times ):
        self.send_command( "consult", self.src )
        assert( self.recv_ok() )
        v1 = int( self.recv_reply() )
        self.send_command( "consult", self.dst )
        assert( self.recv_ok() )
        v2 = int( self.recv_reply() )

        if ( v1 != self.src_base_val - times * 10 ) or \
               ( v2 != self.dst_base_val + times * 10 ):
            raise ValueError( ( "After %s transfers of %s with initial "
                                "values %s and %s, got %s and %s instead " \
                                "of the correct values %s and %s." ) % \
                              ( times, 10, self.src_base_val,
                                self.dst_base_val, v1, v2,
                                self.src_base_val - times * 10,
                                self.dst_base_val + times * 10 ) )
        else:
            return True
    # check()
# StressTestTransfer


class StressTestConsult( StressTest ):
    """Stress Test 'consult' operation."""
    def setup( self ):
        StressTest.setup( self )
        self.send_command( "consult", self.account )
        assert( self.recv_ok() )
        self.base_val = int( self.recv_reply() )

        self.send_command( "deposit", self.account, 10 )
        assert( self.recv_ok() )
    # setup()


    def action( self ):
        self.send_command( "consult", self.account )
        assert( self.recv_ok() )
        v = int( self.recv_reply() )
        try:
            assert( v == 10 + self.base_val )
        except:
            raise SystemExit( ( "consult should return %s, " \
                                "returned %s instead." ) % \
                              ( self.base_val + 10, v ) )
    # action()

    def check( self, times ):
        return True
# StressTestConsult





if __name__ == "__main__":
    tests = {
        "deposit": StressTestDeposit,
        "consult": StressTestConsult,
        "withdraw": StressTestWithdraw,
        "transfer": StressTestTransfer,
        }

    host    = "localhost"
    port    = 8021
    number  = None
    maxtime = 60
    rid     = 0
    test    = None
    verbose = 0


    def usage():
        print """\
Usage:

   %(progname)s [options] <test>

Where options can be:
  -H, --host=HOST        the coordinator hostname or IP.  [default=%(defhost)s]
  -p, --port=PORT        the coordinator port at host.    [default=%(defport)s]
  -n, --number=NUMBER    how many seconds takes to execute NUMBER runs.
  -m, --maxtime=MAXTIME  how many runs it do in MAXTIME seconds.
  -r, --requestid=RID    first request identifier to use  [default=%(defrid)s]
  -v, --verbose=LEVEL    verbosity level.                 [default=%(defverb)s]
  -h, --help             this message.

If no option was given, defaults to -m%(maxtime)s.

test can be one of: %(tests)s.
""" % { "progname": sys.argv[ 0 ],
        "tests": ", ".join( tests.iterkeys() ),
        'maxtime': maxtime,
	"defhost": host,
	"defport": port,
        "defverb": verbose,
        "defrid": rid,
        }
        raise SystemExit
    # usage()

    opts, args = getopt.getopt( sys.argv[ 1: ], "hm:n:H:p:v:r:",
                                [ "help", "maxtime=", "number=",
                                  "host=", "port=", "verbose=",
                                  "requestid=" ] )
    for o, a in opts:
        if o in ( "-n", "--number" ):
            number = int( a )
        elif o in ( "-m", "--maxtime" ):
            maxtime = float( a )
        elif o in ( "-H", "--host" ):
            host = a
        elif o in ( "-p", "--port" ):
            port = int( a )
        elif o in ( "-r", "--requestid" ):
            rid = int( a )
        elif o in ( "-v", "--verbose" ):
            verbose = int( a )
        elif o in ( "-h", "--help" ):
            usage()

    debug.level = verbose

    for arg in args:
        try:
            test = tests[ arg ]
            break
        except KeyError:
            pass
    else:
        raise SystemExit( "You must provide a valid test!" )

    start_time = timer()
    def show_info( signum, frame ):
        t = timer() - start_time
        print "Current request id is: %s. Seconds: %s" % \
              ( st.request_id, t )
    # show_info()
    signal.signal( signal.SIGUSR1, show_info )

    debug( 1, "Current PID is %s" % os.getpid() )
    debug( 2, "You can send 'kill -SIGUSR1 %s' to get progress info." % \
           os.getpid() )


    st = test( ( host, port ) )
    test.set_request_id( rid )
    if number is not None:
        sys.stdout.write( "Running %s times..." % number )
        sys.stdout.flush()
        if debug.level > 0:
            sys.stderr.write( '\n' )
        v = st.run( number )
        sys.stdout.write( ( "%0.2f seconds total " \
                            "(average=%0.4f operations/second)\n" ) % v )

    elif maxtime is not None:
        sys.stdout.write( "How many runs in %s seconds..." % maxtime )
        sys.stdout.flush()
        if debug.level > 0:
            sys.stderr.write( '\n' )
        v = st.howmany( maxtime )
        sys.stdout.write( ( "%s operations (real time=%0.2f seconds, " \
                            "average=%0.4f seconds/operation)\n" ) % v )
# endif
